#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   hoi_pose_traj_estimator_losses.py
@Time    :   2025/03/13 19:34:40
@Author  :   Sanqi Lu
@Contact :   sanqilu376@163.com
@Desc    :   当前文件作用
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from einops import rearrange

def Sinkhorn(cost, reg, num_iterations):
    """
    Sinkhorn算法处理不同点数的点云（N ≠ M）。
    cost: (B, N, M), 预测点云(N)到真实点云(M)的逐对距离
    """
    B, N, M = cost.shape
    K = torch.exp(-cost / reg)
    u = torch.ones(B, N, 1, device=cost.device) / N  # 预测点质量：1/N
    v = torch.ones(B, M, 1, device=cost.device) / M  # 真实点质量：1/M
    
    for _ in range(num_iterations):
        # 交替更新u和v，确保总质量守恒
        v = (1.0 / (torch.bmm(K.transpose(1, 2), u) + 1e-8)) * (1/M)
        u = (1.0 / (torch.bmm(K, v) + 1e-8)) * (1/N)
    
    # 计算传输计划 (B, N, M)
    transport = u * K * v.transpose(1, 2)
    return transport

def Earth_Mover_Distance(pred, gt, reg=0.05, num_iterations=50):
    """
    支持不同点数的EMD实现：
    pred: (B, N, 3), 预测点云（点数N）
    gt: (B, M, 3), 真实点云（点数M）
    """
    B, N, _ = pred.shape
    M = gt.shape[1]
    
    # 计算逐对距离 (B, N, M)
    cost = torch.cdist(pred, gt, p=2.0) ** 2
    
    # 归一化成本矩阵（避免数值溢出）
    cost_normalized = cost / (cost.max(dim=2, keepdim=True)[0] + 1e-8)
    
    # 计算传输计划 (B, N, M)
    transport = Sinkhorn(cost_normalized, reg, num_iterations)
    # print(f"transport: {transport.shape}")
    # 计算EMD：传输计划与原始成本的加权和
    emd = torch.sum(transport * cost, dim=(1, 2))  # (B,)
    return emd

def Chamfer_Distance(pred, gt):
    """
    Computes the Chamfer Distance between two point clouds.
    
    Args:
        pred: torch.Tensor of shape (B, N, 3), predicted point cloud.
        gt: torch.Tensor of shape (B, M, 3), ground truth point cloud.
    
    Returns:
        Mean Chamfer Distance over the batch.
    """
    assert pred.size(0) == gt.size(0), "Batch size mismatch."
    
    # Compute pairwise squared distances
    dist_sq = torch.cdist(pred, gt, p=2.0) ** 2  # (B, N, M)
    # print(f"dist_sq: {dist_sq.shape}")
    # Minimum distance for each predicted point to gt
    min_dist_p_to_g = torch.min(dist_sq, dim=2)[0]  # (B, N)
    # print(f"min_dist_p_to_g: {min_dist_p_to_g.shape}")
    # Minimum distance for each gt point to predicted
    min_dist_g_to_p = torch.min(dist_sq, dim=1)[0]  # (B, M)    
    # Average over points and sum both directions
    chamfer_loss = (torch.mean(min_dist_p_to_g, dim=1) + torch.mean(min_dist_g_to_p, dim=1))
    # print(f"chamfer_loss: {chamfer_loss.shape}")
    return chamfer_loss

def Chamfer_Distance_Blocked(pred, gt, chunk_size=1024):
    B, N, _ = pred.shape
    M = gt.shape[1]
    device = pred.device

    # 分块计算pred到gt的最小距离
    min_dist_p_to_g = torch.full((B, N), float('inf'), device=device)
    for j in range(0, M, chunk_size):
        gt_chunk = gt[:, j:j+chunk_size]
        dist_chunk = torch.cdist(pred, gt_chunk)**2
        min_dist_p_to_g = torch.minimum(min_dist_p_to_g, torch.min(dist_chunk, dim=2).values)

    # 分块计算gt到pred的最小距离
    min_dist_g_to_p = torch.full((B, M), float('inf'), device=device)
    for i in range(0, N, chunk_size):
        pred_chunk = pred[:, i:i+chunk_size]
        dist_chunk = torch.cdist(pred_chunk, gt)**2
        min_dist_g_to_p = torch.minimum(min_dist_g_to_p, torch.min(dist_chunk, dim=1).values)

    chamfer_loss = (torch.mean(min_dist_p_to_g, dim=1) + torch.mean(min_dist_g_to_p, dim=1))
    return chamfer_loss


def compute_hand_loss(pred_hand_traj, gt_hand_traj):
    # [b, t, 2, 21, 3]
    # temporal_weights = self.temporal_weights.to(device=pred_traj.device, dtype=pred_traj.dtype)
    pos_loss = F.smooth_l1_loss(pred_hand_traj, gt_hand_traj, reduction='none') #  * temporal_weights[None, :, None, None]
    pos_loss = pos_loss.mean(dim=(1, 2, 3, 4))

    velocity = pred_hand_traj[:, 1:] - pred_hand_traj[:, :-1]
    gt_velocity = gt_hand_traj[:, 1:] - gt_hand_traj[:, :-1]
    vel_match_loss = F.l1_loss(velocity, gt_velocity, reduction='none')
    vel_match_loss = vel_match_loss.mean(dim=(1, 2, 3, 4))

    acceleration = velocity[:, 1:] - velocity[:, :-1]
    gt_acceleration = gt_velocity[:, 1:] - gt_velocity[:, :-1]
    acc_match_loss = F.l1_loss(acceleration, gt_acceleration, reduction='none')
    acc_match_loss = acc_match_loss.mean(dim=(1, 2, 3, 4))

    return pos_loss, vel_match_loss, acc_match_loss

def compute_obj_loss(pred_obj_traj, gt_obj_traj):
    # [b t m k c] -> [bt, 2*128, 3]
    pred_tool_traj = rearrange(pred_obj_traj[:, :, 0, :, :], "b t k c -> (b t) k c")
    gt_tool_traj = rearrange(gt_obj_traj[:, :, 0, :, :], "b t k c -> (b t) k c")
    tool_chamfer_dist = Chamfer_Distance(pred_tool_traj, gt_tool_traj)

    pred_target_traj = rearrange(pred_obj_traj[:, :, 1, :, :], "b t k c -> (b t) k c")
    gt_target_traj = rearrange(gt_obj_traj[:, :, 1, :, :], "b t k c -> (b t) k c")
    target_chamfer_dist = Chamfer_Distance(pred_target_traj, gt_target_traj)

    chamfer_dist = tool_chamfer_dist + target_chamfer_dist
    # chamfer_dist = Chamfer_Distance_Blocked(pred_obj_traj, gt_obj_traj, chunk_size=256)
    # emd = Earth_Mover_Distance(pred_obj_traj, gt_obj_traj)
    emd = None
    return chamfer_dist, emd

def compute_hoi_pose_traj_loss(pred_traj, gt_traj):
    w_position = 1
    w_velocity = 0.5
    w_acc = 0.2
    w_chamfer = 1
    w_emd = 0.1

    # [b, t, n, c]
    B, T = pred_traj.shape[0], pred_traj.shape[1]
    pred_traj = rearrange(pred_traj, "b t (m k) c -> b t m k c", m=2)
    gt_traj = rearrange(gt_traj, "b t (m k) c -> b t m k c", m=2)

    pred_hand_traj = pred_traj[:, :, :, :21, :]
    gt_hand_traj = gt_traj[:, :, :, :21, :]
    pos_loss, vel_match_loss, acc_match_loss = compute_hand_loss(pred_hand_traj, gt_hand_traj)

    pred_obj_traj = pred_traj[:, :, :, 21:, :]
    gt_obj_traj = gt_traj[:, :, :, 21:, :]
    chamfer_dist, _ = compute_obj_loss(pred_obj_traj, gt_obj_traj)
    chamfer_dist = rearrange(chamfer_dist, "(b t) -> b t", b=B).mean(dim=-1)
    total_loss = w_position * pos_loss + w_velocity * vel_match_loss + w_acc *acc_match_loss + w_chamfer * chamfer_dist # + self.w_emd * emd
    
    return {
        'total_loss': total_loss,
        'pos_loss': pos_loss,
        'vel_match_loss': vel_match_loss,
        'acc_match_loss': acc_match_loss,
        "chamfer_dist": chamfer_dist,
        # "earth_mover_dist": emd,
    }
    